﻿using System;
using System.IO;
using System.Net.WebSockets;
using System.Text;
using System.Threading;

namespace BToolkit
{
    public class WebSocketClient
    {
        public bool isConnecting, isConnected;
        public event Action<string> onSubThreadMessage, onMainThreadMessage;
        public event Action onDisconnect;

        private ClientWebSocket socket;
        private CancellationTokenSource cts;
        private MemoryStream receiveBuffer = new MemoryStream();

        public async CTask<bool> Connect(string url)
        {
            isConnecting = true;
            socket = new ClientWebSocket();
            cts = new CancellationTokenSource();
            try
            {
                await socket.ConnectAsync(new Uri(url), cts.Token);
                ReceiveLoop();
                isConnecting = false;
                isConnected = true;
                return true;
            }
            catch (Exception ex)
            {
                isConnecting = false;
                Console.WriteLine("WebSocket连接失败：" + ex.Message);
            }
            return false;
        }

        public async void Send(string msg)
        {
            if (socket?.State != WebSocketState.Open)
            {
                Console.WriteLine("WebSocket未连接!");
                return;
            }
            try
            {
                var buffer = new ArraySegment<byte>(Encoding.UTF8.GetBytes(msg));
                await socket.SendAsync(buffer, WebSocketMessageType.Text, true, CancellationToken.None);
            }
            catch (Exception ex)
            {
                Console.WriteLine("发送消息失败：" + ex.Message);
            }
        }

        public void Disconnect()
        {
            if (socket?.State == WebSocketState.Open)
            {
                cts.Cancel();
                socket.Dispose();
            }
        }

        private async void ReceiveLoop()
        {
            var buffer = new byte[1024];

            while (socket.State == WebSocketState.Open && !cts.Token.IsCancellationRequested)
            {
                try
                {
                    var receiveResult = await socket.ReceiveAsync(new ArraySegment<byte>(buffer), cts.Token);

                    if (receiveResult.MessageType == WebSocketMessageType.Close)
                    {
                        isConnected = false;
                        MainThread.Run(() => { onDisconnect?.Invoke(); });
                        Console.WriteLine("收到关闭消息，断开连接。");
                        break;
                    }

                    receiveBuffer.Write(buffer, 0, receiveResult.Count);
                    if (receiveResult.EndOfMessage)
                    {
                        byte[] _receiveBuffer = receiveBuffer.ToArray();
                        var msg = Encoding.UTF8.GetString(_receiveBuffer, 0, _receiveBuffer.Length);
                        onSubThreadMessage?.Invoke(msg);
                        MainThread.Run(() => { onMainThreadMessage?.Invoke(msg); });
                        receiveBuffer.SetLength(0);
                        receiveBuffer.Position = 0;
                    }
                }
                catch (Exception ex)
                {
                    isConnected = false;
                    MainThread.Run(() => { onDisconnect?.Invoke(); });
                    Console.WriteLine("接收消息失败：" + ex.Message);
                    break;
                }
            }
        }
    }
}
